import caffe
import caffe.proto.caffe_pb2
import numpy as np
from caffe.io import blobproto_to_array
import os
import logging
import random

FORMAT = '[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
DATEFORMAT = '%H:%M:%S'
logging.basicConfig(level=logging.INFO, format=FORMAT, datefmt=DATEFORMAT)
logging.debug('start')

from tasks import pre_pocess

home_dir = os.getenv('HOME','')

#load the model
net = caffe.Net('bvlc_alexnet/deploy.prototxt',
                'alexnet_snap/caffe_alexnet_train_iter_5000.caffemodel',
                caffe.TEST)


blob = caffe.proto.caffe_pb2.BlobProto()
data = open( 'data_mean.binaryproto' , 'rb' ).read()
blob.ParseFromString(data)
arr = np.array( caffe.io.blobproto_to_array(blob) )


# load input and configure preprocessing
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_mean('data', arr[0].mean(1).mean(1))
transformer.set_transpose('data', (2,0,1))
transformer.set_channel_swap('data', (2,1,0))
transformer.set_raw_scale('data', 255.0)

#note we can change the batch size on-the-fly
#since we classify only one image, we change batch size from 10 to 1
# net.blobs['data'].reshape(1,3,85,85)
#load the image in the data layer

root = home_dir + '/disk/2016年1-10/2/'
flist = os.listdir(root)
for n in flist:
    if n.find('-') >= 0:
        continue
    imgname = root + n
    tmpname = '/tmp/%010d.png' % (random.randint(0, 10000))
    pre_pocess(imgname, tmpname)
    im = caffe.io.load_image(tmpname)
    net.blobs['data'].data[...] = transformer.preprocess('data', im)
    out = net.forward()['prob'][0]
    os.remove(tmpname)
    print(n, out, sep='\t')
